import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Random import get_random_bytes

# 生成32字节 随机密钥
aes_key_file = "key/aes_key.bin"
# os.system(f"openssl rand -out {aes_key_file} 32")


def read_aes_key(file):
    with open(file, "rb") as fp:
        key = fp.read()
        if len(key) in (12, 24, 32):
            return key
        else:
            raise ValueError("Invalid key size")


def aes_encrypt_file(key, file):
    iv = get_random_bytes(16)
    aes_impl = AES.new(key, AES.MODE_CBC, iv)
    with open(file, "rb") as in_fp, open(f"{file}.aes_enc", "wb") as out_fp:
        out_fp.write(iv)
        while True:
            chunk = in_fp.read(64 * 1024)
            if len(chunk) == 0:
                break
            elif len(chunk) % 16 != 0:
                chunk = pad(chunk, AES.block_size)
            encrypt_chunk = aes_impl.encrypt(chunk)
            out_fp.write(encrypt_chunk)


def aes_decrypt_file(key, file):
    with open(file, "rb") as in_fp, open(f"{file}.aes_dec", "wb") as out_fp:
        iv = in_fp.read(16)
        aes_impl = AES.new(key, AES.MODE_CBC, iv)
        while True:
            chunk = in_fp.read(64 * 1024)
            if len(chunk) == 0:
                break
            decrypt_chunk = unpad(aes_impl.decrypt(chunk), AES.block_size)
            out_fp.write(decrypt_chunk)

if __name__ == "__main__":
    key = read_aes_key(aes_key_file)
    file="key/file.txt"
    aes_encrypt_file(key,file)